from typing import Any
from mmdeploy_runtime import Detector
import numpy as np


class MMDetector(object):
    def __init__(
        self, keys, model_dir, threshold=0.5, batch_size=1, device="cuda:0"
    ) -> None:
        self.threshold = threshold
        self.batch_size = batch_size
        self.keys = keys
        if ":" in device:
            device_name, device_id = device.split(":")
            device_id = int(device_id)
        else:
            device_name = device
            device_id = 0
        self.detector = Detector(
            model_path=model_dir, device_name=device_name, device_id=device_id
        )

    def __call__(self, data) -> Any:
        """
        input:
            data: 字典格式或者np.ndarray格式的数据
                注意输入的图片数据格式为[NHWC]
        output:
            data{
                "dets":[
                    [bboxes, score, info] # img 1
                    [bboxes, score, info] # img 2
                    ...

                    ]
            }
        """
        input_tensor = data[self.keys["in_data"]]
        dets = self.detector.batch(input_tensor)
        filter_res = self.filter_box(dets, threshold=self.threshold)
        data[self.keys["out_det"]] = filter_res
        return data

    def release(self):
        pass

    def filter_box(self, result, threshold):
        """根据阈值筛选目标, 并整理格式为 [cls_id, score, x1, y1, x2, y2]"""
        bboxes = []
        for bboxes_, cls_id_, _ in result:
            if bboxes_.shape[0] <= 0:
                bboxes.append(np.empty((0, 6)))
            else:
                bboxes_, score_ = bboxes_[:, 0:4], bboxes_[:, 4]
                a = score_ > threshold
                b = (bboxes_[:,2]-bboxes_[:,0])>0
                c = (bboxes_[:,3]-bboxes_[:,1])>0
                mask = a & b & c
                if not mask.any():
                    bboxes.append(np.empty((0, 6)))
                    continue
                bboxes_ = bboxes_[mask]
                cls_id_ = cls_id_[mask][..., None]
                score_ = score_[mask][..., None]
                bboxes_ = np.concatenate([cls_id_, score_, bboxes_], axis=-1)
                bboxes.append(bboxes_)
        return {"boxes": bboxes}


if __name__ == "__main__":

    def build_labelme_json(bboxes, cls_map,  filename, width, height):
        labelme_json = {
                "version": "5.0.2", 
                "flags": {},
                "shapes": [],
                "imagePath": "..\\data\\"+filename,
                "imageHeight": height,
                "imageWidth": width,
                "imageData":None,
            }
        shapes = []
        for bbox in bboxes:
            cls_id, score, x1,y1,x2,y2 = bbox
            cls_name = cls_map[cls_id]
            shape = {
                "label":cls_name,
                "points":[[x1,y1],[x2,y2]],
                "group_id": None,
                "shape_type": "rectangle",
                "flags":{}
                }
            shapes.append(shape)
        labelme_json["shapes"] = shapes
        return labelme_json

    def det_img_to_labelme():
        import cv2, os, shutil, json
        from glob import glob
        from tqdm import tqdm
        import numpy as np
        dir_path = "/home/smartgis/workspace/data/电力系统/wind_power/预巡/3/det/org/data"
        save_path = "/home/smartgis/workspace/data/电力系统/wind_power/预巡/3/det/org/label"
        if os.path.exists(save_path):
            shutil.rmtree(save_path)
        os.makedirs(save_path)
        cls_map = {
            0:"main",
            1:"other"
        }

        detector = MMDetector(
            model_dir="projects/wind_power/weights/rtmdet_m",
            keys={
                "in_data": "input_data",
                "out_det": "detect",
            },
            device="cuda:0",
        )
        
        file_list = list(glob(dir_path + "/*"))
        for path in tqdm(file_list):
            img_data = cv2.imread(path)
            data = {"input_data": img_data[None, ...]}
            data = detector(data)
            det_result = data["detect"]
            bboxes = det_result["boxes"]
            labelme_json = build_labelme_json(bboxes[0], cls_map,os.path.basename(path),img_data.shape[1],img_data.shape[0])
            with open(save_path+"/{}.json".format(os.path.basename(path).rsplit(".",1)[0]),"w") as f:
                json.dump(labelme_json,f)

    test()
